import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from config import *
from model_advanced import HybridWorldModel

class EnsemblePrediction:
    """集成预测类：结合多种预测策略获得更准确的结果"""
    def __init__(self, model_path="models/hybrid_world_model_best.pth", device=None):
        if device is None:
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        else:
            self.device = device
            
        # 加载模型
        self.hybrid_model = self._load_model(model_path)
        
        # 历史窗口大小
        self.history_length = 5
        
        # 历史状态缓冲区
        self.state_history = []
        
        # 预测信心阈值
        self.confidence_threshold = 0.8
        
        # 规则引擎
        self.rule_engine = RuleBasedPredictionEngine()
        
    def _load_model(self, model_path):
        """加载预训练模型"""
        # 计算特征数
        node_features = self._calculate_node_features()
        edge_features = 4  # 方向特征
        
        # 创建模型
        model = HybridWorldModel(
            node_features=node_features,
            edge_features=edge_features,
            hidden_dim=256,
            history_length=5
        )
        
        # 加载权重
        model.load_state_dict(torch.load(model_path, map_location=self.device))
        model = model.to(self.device)
        model.eval()
        
        return model
    
    def _calculate_node_features(self):
        """计算节点特征数"""
        base_features = CELL_FEATURES
        positional_encoding = 2
        
        # 额外特征
        extra_features = 0
        if USE_DISTANCE_MAPS:
            extra_features += 3
        if USE_DIRECTION_MAPS:
            extra_features += 4
            
        # 路径规划特征
        path_features = 1 + 5 + 1 + 4  # 路径地图 + 动作热图 + 可达性地图 + 目标特征
        
        return base_features + positional_encoding + extra_features + path_features
        
    def predict(self, current_state):
        """预测下一个状态，特别是智能体位置"""
        # 添加当前状态到历史记录
        self.state_history.append(current_state)
        
        # 保持历史记录在固定长度
        if len(self.state_history) > self.history_length:
            self.state_history = self.state_history[-self.history_length:]
            
        # 准备当前状态
        state_tensor = torch.FloatTensor(current_state).unsqueeze(0).to(self.device)
        
        # 准备历史状态
        if len(self.state_history) > 1:
            history_tensor = torch.FloatTensor(np.stack(self.state_history[:-1])).unsqueeze(0).to(self.device)
        else:
            history_tensor = None
            
        # 神经网络预测
        with torch.no_grad():
            env_output, agent_output = self.hybrid_model(state_tensor, history_tensor)
            
        # 提取预测
        agent_pos_pred = agent_output[0, 2:].cpu().numpy()  # 绝对位置(x,y)
        agent_move_pred = agent_output[0, :2].cpu().numpy()  # 相对移动(dx,dy)
        
        # 提取当前智能体位置
        grid_state = self._extract_grid_state(current_state)
        current_agent_pos = self._find_agent_position(grid_state)
        
        if current_agent_pos is None:
            # 如果找不到当前位置，直接使用模型预测
            final_prediction = np.round(agent_pos_pred).astype(int)
        else:
            # 计算基于相对移动的预测
            rel_prediction = np.array(current_agent_pos) + np.round(agent_move_pred).astype(int)
            
            # 计算基于绝对位置的预测
            abs_prediction = np.round(agent_pos_pred).astype(int)
            
            # 让规则引擎也做出预测
            rule_prediction = self.rule_engine.predict(grid_state, current_agent_pos)
            
            # 集成所有预测
            final_prediction = self._ensemble_predictions(
                grid_state, 
                current_agent_pos,
                rel_prediction, 
                abs_prediction, 
                rule_prediction
            )
        
        # 确保预测在有效范围内
        final_prediction[0] = max(0, min(final_prediction[0], ENV_SIZE-1))
        final_prediction[1] = max(0, min(final_prediction[1], ENV_SIZE-1))
        
        return final_prediction
    
    def _extract_grid_state(self, state_features):
        """从特征表示中提取网格状态"""
        # 重塑为网格形式
        features = np.reshape(state_features, (ENV_SIZE, ENV_SIZE, -1))
        
        # 提取基本单元格状态（前CELL_FEATURES个通道）
        cell_states = features[:, :, :CELL_FEATURES]
        
        # 转换为索引形式
        grid_state = np.argmax(cell_states, axis=2)
        
        return grid_state
    
    def _find_agent_position(self, grid_state):
        """在网格状态中找到智能体位置"""
        agent_pos = np.where(grid_state == AGENT)
        
        if len(agent_pos[0]) > 0:
            return (agent_pos[0][0], agent_pos[1][0])
        
        return None
    
    def _ensemble_predictions(self, grid_state, current_pos, rel_pred, abs_pred, rule_pred):
        """集成多种预测方法的结果"""
        # 检查预测是否有效（不在障碍物上）
        predictions = []
        
        # 添加有效预测
        for pred in [rel_pred, abs_pred, rule_pred]:
            i, j = int(pred[0]), int(pred[1])
            if 0 <= i < ENV_SIZE and 0 <= j < ENV_SIZE and grid_state[i, j] != OBSTACLE:
                predictions.append(pred)
        
        if not predictions:
            # 如果没有有效预测，返回当前位置
            return np.array(current_pos)
        
        # 计算预测聚类
        clusters = self._cluster_predictions(predictions)
        
        # 选择最大聚类的中心
        if clusters:
            largest_cluster = max(clusters, key=len)
            final_pred = np.mean(largest_cluster, axis=0)
            return np.round(final_pred).astype(int)
        
        # 如果没有聚类，返回相对预测（通常更可靠）
        return rel_pred
    
    def _cluster_predictions(self, predictions, distance_threshold=1.5):
        """将预测分组到聚类中"""
        if not predictions:
            return []
            
        # 初始化聚类
        clusters = [[predictions[0]]]
        
        # 对每个预测分配到最近的聚类
        for pred in predictions[1:]:
            pred = np.array(pred)
            assigned = False
            
            for cluster in clusters:
                # 计算到聚类中心的距离
                center = np.mean(cluster, axis=0)
                distance = np.linalg.norm(pred - center)
                
                if distance <= distance_threshold:
                    cluster.append(pred)
                    assigned = True
                    break
                    
            if not assigned:
                # 创建新聚类
                clusters.append([pred])
                
        return clusters

class RuleBasedPredictionEngine:
    """基于规则的预测引擎：使用启发式方法预测智能体位置"""
    def __init__(self):
        pass
        
    def predict(self, grid_state, current_pos):
        """根据启发式规则预测下一个智能体位置"""
        if current_pos is None:
            # 如果找不到当前位置，返回零位置
            return np.array([0, 0])
            
        i, j = current_pos
        
        # 找到目标位置
        goal_pos = np.where(grid_state == END)
        if len(goal_pos[0]) > 0:
            goal_i, goal_j = goal_pos[0][0], goal_pos[1][0]
            
            # 计算向目标的方向
            di = goal_i - i
            dj = goal_j - j
            
            # 优先选择曼哈顿距离减少的移动
            possible_moves = []
            
            # 考虑垂直移动
            if di < 0 and i > 0 and grid_state[i-1, j] != OBSTACLE:
                possible_moves.append((-1, 0))  # 上
            elif di > 0 and i < ENV_SIZE-1 and grid_state[i+1, j] != OBSTACLE:
                possible_moves.append((1, 0))   # 下
                
            # 考虑水平移动
            if dj < 0 and j > 0 and grid_state[i, j-1] != OBSTACLE:
                possible_moves.append((0, -1))  # 左
            elif dj > 0 and j < ENV_SIZE-1 and grid_state[i, j+1] != OBSTACLE:
                possible_moves.append((0, 1))   # 右
                
            # 如果有可能的移动，选择最佳的
            if possible_moves:
                # 优先选择距离目标最近的移动
                best_move = min(possible_moves, key=lambda m: abs(goal_i-(i+m[0])) + abs(goal_j-(j+m[1])))
                return np.array([i + best_move[0], j + best_move[1]])
        
        # 如果没有目标或无法向目标移动，探索周围空间
        for di, dj in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
            ni, nj = i + di, j + dj
            if 0 <= ni < ENV_SIZE and 0 <= nj < ENV_SIZE and grid_state[ni, nj] != OBSTACLE:
                return np.array([ni, nj])
                
        # 如果无法移动，保持原位
        return np.array([i, j])